//
//  CollectionFlexLayout.swift
//  Risac
//
//  Created by remy on 2019/11/11.
//  Copyright © 2019 com.Risac. All rights reserved.
//

import UIKit

open class CollectionFlexLayout: UICollectionViewLayout {
    public weak var model: CollectionModel?
    public let adjustSize: Bool
    public let direction: UICollectionView.ScrollDirection
    private var cellAttrs: [IndexPath: UICollectionViewLayoutAttributes] = [:]
    private var sectionAttrs: [Int: UICollectionViewLayoutAttributes] = [:]
    private var contentSize: CGSize = .zero
    override open var collectionViewContentSize: CGSize {
        return self.contentSize
    }
    
    public init(model: CollectionModel,
                adjustSize: Bool = false,
                direction: UICollectionView.ScrollDirection = .vertical) {
        self.model = model
        self.adjustSize = adjustSize
        self.direction = direction
        super.init()
    }
    
    required public init?(coder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }
    
    override open func prepare() {
        super.prepare()
        prepareAttributes()
    }
    
    private func prepareAttributes() {
        guard let model = model else { return }
        guard let collectionView = self.collectionView else { return }
        sectionAttrs.removeAll()
        cellAttrs.removeAll()
        let viewSize: CGSize = collectionView.size
        var maxCellLimit: CGPoint = .zero
        var offset: CGPoint = .zero
        if direction == .vertical {
            model.sectionItems.forEachEnumerated { (sIndex, section) in
                let indexPath = IndexPath(row: 0, section: sIndex)
                let sectionMargin = section.sectionMargin
                let sectionInset = section.sectionInset
                let lineSpacing = section.lineSpacing
                let itemSpacing = section.itemSpacing
                let attr = UICollectionViewLayoutAttributes(forSupplementaryViewOfKind: UICollectionView.elementKindSectionStage, with: indexPath)
                attr.zIndex = -1
                sectionAttrs[sIndex] = attr
                let sectionStartPoint = CGPoint(x: sectionMargin.left, y: (sIndex > 0 ? maxCellLimit.y + section.sectionSpacing : 0.0) + sectionMargin.top)
                var last: (UICollectionViewLayoutAttributes, CollectionCellItem)?
                let baseCellLeft = sectionMargin.left + sectionInset.left
                var baseCellTop = sectionStartPoint.y + sectionInset.top
                var maxCellRight = viewSize.width - sectionInset.right - sectionMargin.right
                if adjustSize {
                    maxCellRight = section.cellItems.reduce(maxCellRight, { max($0, baseCellLeft + $1.cellSize.width) })
                }
                section.cellItems.forEachEnumerated { (cIndex, cell) in
                    let indexPath = IndexPath(row: cIndex, section: sIndex)
                    let cellSize = cell.cellSize
                    let cellMargin = cell.cellMargin
                    if let last = last {
                        offset.x = last.0.frame.maxX + last.1.cellMargin.right + itemSpacing + cellMargin.left
                        if offset.x + cellSize.width + cellMargin.right > maxCellRight {
                            offset.x = baseCellLeft + cellMargin.left
                            baseCellTop = maxCellLimit.y + lineSpacing
                        }
                    } else {
                        offset.x = baseCellLeft + cellMargin.left
                    }
                    offset.y = baseCellTop + cellMargin.top
                    maxCellLimit.x = max(maxCellLimit.x, offset.x + cellSize.width + cellMargin.right);
                    maxCellLimit.y = max(maxCellLimit.y, offset.y + cellSize.height + cellMargin.bottom)
                    let attr = UICollectionViewLayoutAttributes(forCellWith: indexPath)
                    attr.frame = CGRect(origin: offset, size: cellSize)
                    cellAttrs[indexPath] = attr
                    last = (attr, cell)
                }
                maxCellLimit.y += sectionInset.bottom + sectionMargin.bottom
                if adjustSize {
                    attr.frame = CGRect(origin: sectionStartPoint, size: CGSize(width: maxCellLimit.x - sectionMargin.left + sectionInset.right, height: maxCellLimit.y - sectionStartPoint.y - sectionMargin.bottom))
                    contentSize = CGSize(width: max(contentSize.width, maxCellLimit.x + sectionInset.right + sectionMargin.right), height: maxCellLimit.y)
                } else {
                    attr.frame = CGRect(origin: sectionStartPoint, size: CGSize(width: viewSize.width - sectionMargin.left - sectionMargin.right, height: maxCellLimit.y - sectionStartPoint.y - sectionMargin.bottom))
                    contentSize = CGSize(width: viewSize.width, height: maxCellLimit.y)
                }
                maxCellLimit.x = 0.0
            }
        } else {
            model.sectionItems.forEachEnumerated { (sIndex, section) in
                let indexPath = IndexPath(row: 0, section: sIndex)
                let sectionMargin = section.sectionMargin
                let sectionInset = section.sectionInset
                let itemSpacing = section.itemSpacing
                let attr = UICollectionViewLayoutAttributes(forSupplementaryViewOfKind: UICollectionView.elementKindSectionStage, with: indexPath)
                attr.zIndex = -1
                sectionAttrs[sIndex] = attr
                let sectionStartPoint = CGPoint(x: (sIndex > 0 ? maxCellLimit.x + section.sectionSpacing : 0.0) + sectionMargin.left, y: sectionMargin.top)
                var last: (UICollectionViewLayoutAttributes, CollectionCellItem)?
                let baseCellLeft = sectionStartPoint.x + sectionInset.left
                let baseCellTop = sectionMargin.top + sectionInset.top
                section.cellItems.forEachEnumerated { (cIndex, cell) in
                    let indexPath = IndexPath(row: cIndex, section: sIndex)
                    let cellSize = cell.cellSize
                    let cellMargin = cell.cellMargin
                    if let last = last {
                        offset.x = last.0.frame.maxX + last.1.cellMargin.right + itemSpacing + cellMargin.left
                    } else {
                        offset.x = baseCellLeft + cellMargin.left
                    }
                    offset.y = baseCellTop + cellMargin.top
                    maxCellLimit.x = offset.x + cellSize.width + cellMargin.right
                    maxCellLimit.y = max(maxCellLimit.y, offset.y + cellSize.height + cellMargin.bottom);
                    let attr = UICollectionViewLayoutAttributes(forCellWith: indexPath)
                    attr.frame = CGRect(origin: offset, size: cellSize)
                    cellAttrs[indexPath] = attr
                    last = (attr, cell)
                }
                maxCellLimit.x += sectionInset.right + sectionMargin.right
                if adjustSize {
                    attr.frame = CGRect(origin: sectionStartPoint, size: CGSize(width: maxCellLimit.x - sectionStartPoint.x - sectionMargin.right, height: maxCellLimit.y - sectionMargin.top + sectionInset.bottom))
                    contentSize = CGSize(width: maxCellLimit.x, height: max(contentSize.height, maxCellLimit.y + sectionInset.bottom + sectionMargin.bottom))
                } else {
                    attr.frame = CGRect(origin: sectionStartPoint, size: CGSize(width: maxCellLimit.x - sectionStartPoint.x - sectionMargin.right, height: viewSize.height - sectionMargin.top - sectionMargin.bottom))
                    contentSize = CGSize(width: maxCellLimit.x, height: viewSize.height)
                }
                maxCellLimit.y = 0.0
            }
        }
    }
    
    override open func layoutAttributesForItem(at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        return cellAttrs[indexPath]
    }
    
    override open func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        return cellAttrs.values.filter { $0.frame.intersects(rect) } + sectionAttrs.values.filter { $0.frame.intersects(rect) }
    }
    
    override open func layoutAttributesForSupplementaryView(ofKind elementKind: String, at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        if elementKind == UICollectionView.elementKindSectionStage {
            return sectionAttrs[indexPath.section]
        }
        return super.layoutAttributesForSupplementaryView(ofKind: elementKind, at: indexPath)
    }
}
